package first.group.ssh.keyexchange;

import first.group.ssh.ConnectionCallback;
import first.group.ssh.hostkey.HostKeyAlgorithm;
import first.group.ssh.packet.DiffieHellmanGroupExchangeGroupPacket;
import first.group.ssh.packet.DiffieHellmanGroupExchangeInitPacket;
import first.group.ssh.packet.DiffieHellmanGroupExchangeReplyPacket;
import first.group.ssh.packet.DiffieHellmanGroupExchangeRequestPacket;
import first.group.ssh.packet.KeyExchangeInitPacket;
import first.group.ssh.packet.Packet;
import first.group.ssh.packet.writer.BigIntegerWriter;
import first.group.ssh.packet.writer.StringWriter;
import first.group.ssh.util.BigIntegerUtil;
import first.group.ssh.version.Version;
import java.io.IOException;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Random;
import java.util.logging.Level;
import java.util.logging.Logger;

public final class DiffieHellmanGroupExchangeSha256 extends KeyExchangeMethod {

    private static final int MIN = 512, N = 1024, MAX = 4096;

    private static final Logger logger = Logger.getLogger("DiffieHellmanGroupExchangeSha256");

    @Override
    public String getName() {
        return "diffie-hellman-group-exchange-sha256";
    }

    @Override
    public MessageDigest newDigest() {
        try {
            return MessageDigest.getInstance("SHA-256");
        } catch (NoSuchAlgorithmException ex) {
            throw new RuntimeException("Digest SHA-256 is not avaliable", ex);
        }
    }

    @Override
    public Exchanger newInstance(
            Version clientVersion,
            Version serverVersion,
            KeyExchangeInitPacket clientInitPacket,
            KeyExchangeInitPacket serverInitPacket,
            Random random,
            ConnectionCallback connectionCallback,
            KeyExchangeMethod.Callback exchangeCallback
    ) {
        return new Exchanger(
                clientVersion,
                serverVersion,
                clientInitPacket,
                serverInitPacket,
                random,
                connectionCallback,
                exchangeCallback);
    }

    private static class Exchanger extends KeyExchangeMethod.Exchanger {

        private enum Status {
            WAIT_GROUP,
            WAIT_REPLY
        }
        private Status status;
        private final Version clientVersion, serverVersion;
        private final KeyExchangeInitPacket clientInitPacket, serverInitPacket;
        private final ConnectionCallback connectionCallback;
        private final KeyExchangeMethod.Callback exchangeCallback;
        private final Random random;

        private BigInteger p, g, x, e, f, k;

        public Exchanger(
                Version clientVersion,
                Version serverVersion,
                KeyExchangeInitPacket clientInitPacket,
                KeyExchangeInitPacket serverInitPacket,
                Random random,
                ConnectionCallback connectionCallback,
                KeyExchangeMethod.Callback exchangeCallback
        ) {
            this.clientVersion = clientVersion;
            this.serverVersion = serverVersion;
            this.clientInitPacket = clientInitPacket;
            this.serverInitPacket = serverInitPacket;
            this.random = random;
            this.connectionCallback = connectionCallback;
            this.exchangeCallback = exchangeCallback;
        }

        @Override
        public void init() {
            this.status = Status.WAIT_GROUP;
            sendRequestPacket();
        }

        private void sendRequestPacket() {
            connectionCallback.appendPacket(new DiffieHellmanGroupExchangeRequestPacket(MIN, N, MAX));
        }

        private void sendInitPacket() {
            var limit = p.subtract(BigInteger.ONE).divide(BigInteger.TWO);
            x = BigIntegerUtil.randomBigInteger(random, BigInteger.ONE, limit);
            e = g.modPow(x, p);

            connectionCallback.appendPacket(new DiffieHellmanGroupExchangeInitPacket(e));
        }

        private byte[] generateHash(DiffieHellmanGroupExchangeReplyPacket packet) {
            var buffer = ByteBuffer.allocate(32768);
            buffer.order(ByteOrder.BIG_ENDIAN);

            StringWriter.writeString(buffer, clientVersion.toString());
            StringWriter.writeString(buffer, serverVersion.toString());
            StringWriter.writePacket(buffer, clientInitPacket);
            StringWriter.writePacket(buffer, serverInitPacket);

            packet.hostKey.write(buffer);

            buffer.putInt(MIN);
            buffer.putInt(N);
            buffer.putInt(MAX);
            BigIntegerWriter.writeBigInteger(buffer, p);
            BigIntegerWriter.writeBigInteger(buffer, g);
            BigIntegerWriter.writeBigInteger(buffer, e);
            BigIntegerWriter.writeBigInteger(buffer, f);
            BigIntegerWriter.writeBigInteger(buffer, k);

            buffer.flip();
            try {
                var digest = MessageDigest.getInstance("SHA-256");
                digest.update(buffer);
                return digest.digest();
            } catch (NoSuchAlgorithmException ex) {
                throw new RuntimeException("Digest failed", ex);
            }
        }

        @SuppressWarnings("unchecked")
        @Override
        public void read(Packet packet) throws IOException {
            switch (status) {
                case WAIT_GROUP -> {
                    var groupPacket = (DiffieHellmanGroupExchangeGroupPacket) packet;
                    logger.log(Level.FINER, "Received group: {0}", groupPacket);
                    p = groupPacket.p;
                    g = groupPacket.g;
                    status = Status.WAIT_REPLY;
                    sendInitPacket();
                }
                case WAIT_REPLY -> {
                    var replyPacket = (DiffieHellmanGroupExchangeReplyPacket) packet;
                    logger.log(Level.FINER, "Received reply: {0}", replyPacket);
                    f = replyPacket.f;

                    k = f.modPow(x, p);
                    if (k.compareTo(BigInteger.ONE) < 0 || k.compareTo(p.subtract(BigInteger.ONE)) > 0) {
                        throw new IOException("k not in range (1, p - 1)");
                    }

                    var hash = generateHash(replyPacket);
                    var hostKey = (HostKeyAlgorithm.HostKey<HostKeyAlgorithm>) replyPacket.hostKey;
                    var hostSignature = (HostKeyAlgorithm.HostSignature<HostKeyAlgorithm>) replyPacket.hostSignature;
                    if (!hostKey.check(hash, hostSignature)) {
                        throw new IOException("Signature check failed!");
                    }

                    exchangeCallback.onKeyExchanged(k, hash);
                }
                default -> throw new AssertionError(status.name());
            }
        }

        @Override
        public boolean shouldRead() {
            return true;
        }
    }
}
